#include <cuda_runtime.h>

// ============================================================
// FP64→FP32精度转换kernel (double to float)
// ============================================================
__global__ void convertDoubleToFloatKernel(
    int n,
    const double* __restrict__ d_input,
    float* __restrict__ d_output
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        d_output[idx] = static_cast<float>(d_input[idx]);
    }
}

// ============================================================
// FP32→FP64精度转换kernel (float to double)
// ============================================================
__global__ void convertFloatToDoubleKernel(
    int n,
    const float* __restrict__ d_input,
    double* __restrict__ d_output
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        d_output[idx] = static_cast<double>(d_input[idx]);
    }
}

// Jacobi预条件: z = M^(-1) * r (FP64版本)
// M^(-1)是对角矩阵的逆（即element-wise乘法）
__global__ void jacobiPreconditionKernel(
    int n,
    const double* __restrict__ d_M_inv,  // 对角逆
    const double* __restrict__ d_r,      // 输入残差向量
    double* __restrict__ d_z             // 输出预条件残差
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        d_z[idx] = d_M_inv[idx] * d_r[idx];
    }
}

// Jacobi预条件: z = M^(-1) * r (FP32版本)
__global__ void jacobiPreconditionKernel_fp32(
    int n,
    const float* __restrict__ d_M_inv,   // 对角逆（FP32）
    const float* __restrict__ d_r,       // 输入残差向量（FP32）
    float* __restrict__ d_z              // 输出预条件残差（FP32）
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < n) {
        d_z[idx] = d_M_inv[idx] * d_r[idx];
    }
}

// C接口
extern "C" {
    // FP64→FP32转换
    void convertDoubleToFloat(int n, const double* d_input, float* d_output) {
        int blockSize = 256;
        int numBlocks = (n + blockSize - 1) / blockSize;
        convertDoubleToFloatKernel<<<numBlocks, blockSize>>>(n, d_input, d_output);
        cudaDeviceSynchronize();
    }

    // FP32→FP64转换
    void convertFloatToDouble(int n, const float* d_input, double* d_output) {
        int blockSize = 256;
        int numBlocks = (n + blockSize - 1) / blockSize;
        convertFloatToDoubleKernel<<<numBlocks, blockSize>>>(n, d_input, d_output);
        cudaDeviceSynchronize();
    }

    // Jacobi预条件（FP64版本）
    void jacobiPrecondition(
        int n,
        const double* d_M_inv,
        const double* d_r,
        double* d_z
    ) {
        int blockSize = 256;
        int numBlocks = (n + blockSize - 1) / blockSize;

        jacobiPreconditionKernel<<<numBlocks, blockSize>>>(
            n, d_M_inv, d_r, d_z
        );

        // 等待kernel完成
        cudaDeviceSynchronize();
    }

    // Jacobi预条件（FP32版本）
    void jacobiPrecondition_fp32(
        int n,
        const float* d_M_inv,
        const float* d_r,
        float* d_z
    ) {
        int blockSize = 256;
        int numBlocks = (n + blockSize - 1) / blockSize;

        jacobiPreconditionKernel_fp32<<<numBlocks, blockSize>>>(
            n, d_M_inv, d_r, d_z
        );

        cudaDeviceSynchronize();
    }
}
